const INIT_D: &[u16] = &[0x44d7, 0x26bc, 0x626b, 0x135e, 0x5789, 0x35e2, 0x7135, 0x09af, 0x4d78, 0x2f13, 0x6bc4, 0x1af1, 0x5e26, 0x3c4d, 0x789a, 0x47ac];
const S0: &[&[u8]] = &[
    &[0x3e, 0x72, 0x5b, 0x47, 0xca, 0xe0, 0x00, 0x33, 0x04, 0xd1, 0x54, 0x98, 0x09, 0xb9, 0x6d, 0xcb],
    &[0x7b, 0x1b, 0xf9, 0x32, 0xaf, 0x9d, 0x6a, 0xa5, 0xb8, 0x2d, 0xfc, 0x1d, 0x08, 0x53, 0x03, 0x90],
    &[0x4d, 0x4e, 0x84, 0x99, 0xe4, 0xce, 0xd9, 0x91, 0xdd, 0xb6, 0x85, 0x48, 0x8b, 0x29, 0x6e, 0xac],
    &[0xcd, 0xc1, 0xf8, 0x1e, 0x73, 0x43, 0x69, 0xc6, 0xb5, 0xbd, 0xfd, 0x39, 0x63, 0x20, 0xd4, 0x38],
    &[0x76, 0x7d, 0xb2, 0xa7, 0xcf, 0xed, 0x57, 0xc5, 0xf3, 0x2c, 0xbb, 0x14, 0x21, 0x06, 0x55, 0x9b],
    &[0xe3, 0xef, 0x5e, 0x31, 0x4f, 0x7f, 0x5a, 0xa4, 0x0d, 0x82, 0x51, 0x49, 0x5f, 0xba, 0x58, 0x1c],
    &[0x4a, 0x16, 0xd5, 0x17, 0xa8, 0x92, 0x24, 0x1f, 0x8c, 0xff, 0xd8, 0xae, 0x2e, 0x01, 0xd3, 0xad],
    &[0x3b, 0x4b, 0xda, 0x46, 0xeb, 0xc9, 0xde, 0x9a, 0x8f, 0x87, 0xd7, 0x3a, 0x80, 0x6f, 0x2f, 0xc8],
    &[0xb1, 0xb4, 0x37, 0xf7, 0x0a, 0x22, 0x13, 0x28, 0x7c, 0xcc, 0x3c, 0x89, 0xc7, 0xc3, 0x96, 0x56],
    &[0x07, 0xbf, 0x7e, 0xf0, 0x0b, 0x2b, 0x97, 0x52, 0x35, 0x41, 0x79, 0x61, 0xa6, 0x4c, 0x10, 0xfe],
    &[0xbc, 0x26, 0x95, 0x88, 0x8a, 0xb0, 0xa3, 0xfb, 0xc0, 0x18, 0x94, 0xf2, 0xe1, 0xe5, 0xe9, 0x5d],
    &[0xd0, 0xdc, 0x11, 0x66, 0x64, 0x5c, 0xec, 0x59, 0x42, 0x75, 0x12, 0xf5, 0x74, 0x9c, 0xaa, 0x23],
    &[0x0e, 0x86, 0xab, 0xbe, 0x2a, 0x02, 0xe7, 0x67, 0xe6, 0x44, 0xa2, 0x6c, 0xc2, 0x93, 0x9f, 0xf1],
    &[0xf6, 0xfa, 0x36, 0xd2, 0x50, 0x68, 0x9e, 0x62, 0x71, 0x15, 0x3d, 0xd6, 0x40, 0xc4, 0xe2, 0x0f],
    &[0x8e, 0x83, 0x77, 0x6b, 0x25, 0x05, 0x3f, 0x0c, 0x30, 0xea, 0x70, 0xb7, 0xa1, 0xe8, 0xa9, 0x65],
    &[0x8d, 0x27, 0x1a, 0xdb, 0x81, 0xb3, 0xa0, 0xf4, 0x45, 0x7a, 0x19, 0xdf, 0xee, 0x78, 0x34, 0x60]
];

const S1: &[&[u8]] = &[
    &[0x55, 0xc2, 0x63, 0x71, 0x3b, 0xc8, 0x47, 0x86, 0x9f, 0x3c, 0xda, 0x5b, 0x29, 0xaa, 0xfd, 0x77],
    &[0x8c, 0xc5, 0x94, 0x0c, 0xa6, 0x1a, 0x13, 0x00, 0xe3, 0xa8, 0x16, 0x72, 0x40, 0xf9, 0xf8, 0x42],
    &[0x44, 0x26, 0x68, 0x96, 0x81, 0xd9, 0x45, 0x3e, 0x10, 0x76, 0xc6, 0xa7, 0x8b, 0x39, 0x43, 0xe1],
    &[0x3a, 0xb5, 0x56, 0x2a, 0xc0, 0x6d, 0xb3, 0x05, 0x22, 0x66, 0xbf, 0xdc, 0x0b, 0xfa, 0x62, 0x48],
    &[0xdd, 0x20, 0x11, 0x06, 0x36, 0xc9, 0xc1, 0xcf, 0xf6, 0x27, 0x52, 0xbb, 0x69, 0xf5, 0xd4, 0x87],
    &[0x7f, 0x84, 0x4c, 0xd2, 0x9c, 0x57, 0xa4, 0xbc, 0x4f, 0x9a, 0xdf, 0xfe, 0xd6, 0x8d, 0x7a, 0xeb],
    &[0x2b, 0x53, 0xd8, 0x5c, 0xa1, 0x14, 0x17, 0xfb, 0x23, 0xd5, 0x7d, 0x30, 0x67, 0x73, 0x08, 0x09],
    &[0xee, 0xb7, 0x70, 0x3f, 0x61, 0xb2, 0x19, 0x8e, 0x4e, 0xe5, 0x4b, 0x93, 0x8f, 0x5d, 0xdb, 0xa9],
    &[0xad, 0xf1, 0xae, 0x2e, 0xcb, 0x0d, 0xfc, 0xf4, 0x2d, 0x46, 0x6e, 0x1d, 0x97, 0xe8, 0xd1, 0xe9],
    &[0x4d, 0x37, 0xa5, 0x75, 0x5e, 0x83, 0x9e, 0xab, 0x82, 0x9d, 0xb9, 0x1c, 0xe0, 0xcd, 0x49, 0x89],
    &[0x01, 0xb6, 0xbd, 0x58, 0x24, 0xa2, 0x5f, 0x38, 0x78, 0x99, 0x15, 0x90, 0x50, 0xb8, 0x95, 0xe4],
    &[0xd0, 0x91, 0xc7, 0xce, 0xed, 0x0f, 0xb4, 0x6f, 0xa0, 0xcc, 0xf0, 0x02, 0x4a, 0x79, 0xc3, 0xde],
    &[0xa3, 0xef, 0xea, 0x51, 0xe6, 0x6b, 0x18, 0xec, 0x1b, 0x2c, 0x80, 0xf7, 0x74, 0xe7, 0xff, 0x21],
    &[0x5a, 0x6a, 0x54, 0x1e, 0x41, 0x31, 0x92, 0x35, 0xc4, 0x33, 0x07, 0x0a, 0xba, 0x7e, 0x0e, 0x34],
    &[0x88, 0xb1, 0x98, 0x7c, 0xf3, 0x3d, 0x60, 0x6c, 0x7b, 0xca, 0xd3, 0x1f, 0x32, 0x65, 0x04, 0x28],
    &[0x64, 0xbe, 0x85, 0x9b, 0x2f, 0x59, 0x8a, 0xd7, 0xb0, 0x25, 0xac, 0xaf, 0x12, 0x03, 0xe2, 0xf2]
];

pub struct Zuc {
    key: Vec<u8>,
    init_vec: Vec<u8>,
    s: Vec<u32>,
    r1: u32,
    r2: u32,
}

impl Default for Zuc {
    fn default() -> Self {
        Self {
            key: vec![0u8; 16],
            init_vec: vec![0u8; 16],
            s: vec![0u32; 16],
            r1: 0,
            r2: 0,
        }
    }
}

impl Zuc {
    pub fn new(key: Vec<u8>, init_vec: Vec<u8>) -> Self {
        Self {
            key,
            init_vec,
            s: vec![0u32; 16],
            r1: 0,
            r2: 0,
        }
    }

    pub fn init(&mut self) {
        for i in 0..16 {
            self.s[i] = ((self.key[i] as u32) << 23) | ((INIT_D[i] as u32) << 8) | (self.init_vec[i] as u32);
        }
        for _ in 0..32 {
            let x = self.br();
            let w = self._f(&x);
            self.init_mode(w >> 1);
            // println!("{:0x}, {:0x}, {:0x}, {:0x},{:0x},{:0x},{:0x},{:0x}", x[0], x[1], x[2], x[3], self.r1, self.r2, w, self.s[15]);
        }
        // println!("{:0x}, {:0x}", self.r1, self.r2);
    }

    pub fn work(&mut self, encode_length: u32) -> Vec<u32> {
        let mut res = vec![];
        let x = self.br();
        let _ = self._f(&x);
        self.work_mode();
        for _ in 0..encode_length {
            res.push(self.gen());
        }
        res
    }

    fn gen(&mut self) -> u32 {
        let x = self.br();
        let w = self._f(&x);
        let z = w ^ x[3];
        self.work_mode();
        z
    }

    fn br(&self) -> Vec<u32> {
        let mut x = vec![0u32; 4];
        x[0] = _concat(high_16bit(self.s[15], 31), low_16bit(self.s[14]));
        x[1] = _concat(low_16bit(self.s[11]), high_16bit(self.s[9], 31));
        x[2] = _concat(low_16bit(self.s[7]), high_16bit(self.s[5], 31));
        x[3] = _concat(low_16bit(self.s[2]), high_16bit(self.s[0], 31));
        x
    }

    fn _f(&mut self, x: &Vec<u32>) -> u32 {
        let w = add_mod32(x[0] ^ self.r1, self.r2);
        let w1 = add_mod32(self.r1, x[1]);
        let w2 = self.r2 ^ x[2];
        self.r1 = self._s(_l1(_concat(low_16bit(w1), high_16bit(w2, 32))));
        self.r2 = self._s(_l2(_concat(low_16bit(w2), high_16bit(w1, 32))));
        w
    }

    fn _s(&self, a: u32) -> u32 {
        let mut bytes = [0u8; 4];
        let mut s_bytes = [0u8; 4];
        bytes[0] = ((a >> 24) & 0xff) as u8;
        bytes[1] = ((a >> 16) & 0xff) as u8;
        bytes[2] = ((a >> 8) & 0xff) as u8;
        bytes[3] = (a & 0xff) as u8;
        for (i, byte) in bytes.iter().enumerate() {
            let r = byte >> 4;
            let c = byte & 0xf;
            s_bytes[i] = match i == 0 || i == 2 {
                true => S0[r as usize][c as usize],
                false => S1[r as usize][c as usize],
            };
        }
        ((s_bytes[0] as u32) << 24) | ((s_bytes[1] as u32) << 16) | ((s_bytes[2] as u32) << 8) | (s_bytes[3] as u32)
    }

    fn init_mode(&mut self, u: u32) {
        let v = ((1u64 << 15) * (self.s[15] as u64)) + ((1u64 << 17) * (self.s[13] as u64)) + ((1u64 << 21) * (self.s[10] as u64)) + ((1u64 << 20) * (self.s[4] as u64)) + (((1u64 << 8) + 1) * (self.s[0] as u64));
        let v = v % ((1 << 31) - 1);
        let v = ((v + (u as u64)) % ((1 << 31) - 1)) as u32;
        let v = match v {
            0 => (1 << 31) - 1,
            _ => v,
        };
        let mut ss = vec![0; 16];
        ss[0..15].copy_from_slice(&self.s[1..16]);
        ss[15] = v;
        self.s = ss;
    }

    fn work_mode(&mut self) {
        let v = ((1u64 << 15) * (self.s[15] as u64)) + ((1u64 << 17) * (self.s[13] as u64)) + ((1u64 << 21) * (self.s[10] as u64)) + ((1u64 << 20) * (self.s[4] as u64)) + (((1u64 << 8) + 1) * (self.s[0] as u64));
        let v = v % ((1 << 31) - 1);
        let v = match v {
            0 => (1 << 31) - 1,
            _ => v as u32,
        };
        let mut ss = vec![0; 16];
        ss[0..15].copy_from_slice(&self.s[1..16]);
        ss[15] = v;
        self.s = ss;
    }
}

fn high_16bit(a: u32, bits: u32) -> u32 {
    (a >> (16 - (32 - bits))) & 0xffff
}

fn low_16bit(a: u32) -> u32 {
    a & 0xffff
}

fn _concat(a: u32, b: u32) -> u32 {
    (a << 16) | b
}

fn add_mod32(a: u32, b: u32) -> u32 {
    (((a as u64) + (b as u64)) & 0xffffffff) as u32
}

fn rotate_left(val: u32, n: u32) -> u32 {
    (val << n) | (val >> (32 - n))
}

fn _l1(a: u32) -> u32 {
    // 2, 10 ,18, 24
    a ^ rotate_left(a, 2) ^ rotate_left(a, 10) ^ rotate_left(a, 18) ^ rotate_left(a, 24)
}

fn _l2(a: u32) -> u32 {
    // 8, 14, 22,30
    a ^ rotate_left(a, 8) ^ rotate_left(a, 14) ^ rotate_left(a, 22) ^ rotate_left(a, 30)
}